from langchain_core.callbacks import StdOutCallbackHandler
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, trim_messages
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.chat_history import InMemoryChatMessageHistory, BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

model = ChatOpenAI(model='Qwen/Qwen2.5-7B-Instruct').bind(
    logprobs=True)
# 消息历史存储
store = {}


# 获取会话历史记录
def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]


# 创建带有 Prompt 模板的链
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant. Answer all questions in {language}."),
    MessagesPlaceholder(variable_name="messages")
])

# 将消息修剪工具添加到链中，避免消息过多超出上下文窗口
trimmer = trim_messages(
    token_counter=len,
    max_tokens=5,  # 限制最大 token 数
    strategy="last",  # 保留最后的消息
    # token_counter=model,  # 使用模型计算 token
    include_system=True,
    allow_partial=False,
    start_on="human"  # 从人类消息开始修剪
)

# 定义处理消息历史和修剪功能的链
from operator import itemgetter
from langchain_core.runnables import RunnablePassthrough


chain = (
        RunnablePassthrough.assign(messages=itemgetter("messages") | trimmer)  # 先修剪消息
        | prompt  # 然后传递到 Prompt 模板
        | model  # 调用模型生成响应
)

# 将消息历史包装到链中
with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="messages"
)

# 测试会话：输入用户消息，并指定语言为西班牙语
config = {"configurable": {"session_id": "10086"}}

# 初始消息

response = with_message_history.invoke(
    {"messages": [HumanMessage(content="Hi! 我是小明.")], "language": "chinese"},
    config=config,
)
print("Response:", response.content)

# 继续对话
response = with_message_history.invoke(
    {"messages": [HumanMessage(content="我是谁?")], "language": "Spanish"},
    config=config,
)
print("Response:", response.content)

# 流式响应
print("\nStreaming Response:")
for r in with_message_history.stream(
        {
            "messages": [HumanMessage(content="讲个笑话.")],
            "language": "English"
        },
        config=config,
):
    print(r.content, end="")

print('end')
